# Copyright (c) OpenMMLab. All rights reserved.
import random
import warnings

import mmcv
import numpy as np
import torch
import torch.distributed as dist
from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
from mmcv.runner import HOOKS, build_optimizer, build_runner, get_dist_info, OptimizerHook
from mmcv.utils import build_from_cfg

from mmseg import digit_version
from mmseg.core import DistEvalHook, EvalHook
from mmseg.datasets import build_dataloader, build_dataset
from mmseg.utils import find_latest_checkpoint, get_root_logger, dataloader_visulization

## added by LYU: 2022/03/28
from mmseg.core.optimizer import build_optimizers as build_optimizers_multi
## added by LYU: 2022/04/06
from mmseg.core.ddp_wrapper import DistributedDataParallelWrapper
from mmcv.runner import IterBasedRunner

def init_random_seed(seed=None, device='cuda'):
    """Initialize random seed.

    If the seed is not set, the seed will be automatically randomized,
    and then broadcast to all processes to prevent some potential bugs.
    Args:
        seed (int, Optional): The seed. Default to None.
        device (str): The device where the seed will be put on.
            Default to 'cuda'.
    Returns:
        int: Seed to be used.
    """
    if seed is not None:
        return seed

    # Make sure all ranks share the same random seed to prevent
    # some potential bugs. Please refer to
    # https://github.com/open-mmlab/mmdetection/issues/6339
    rank, world_size = get_dist_info()
    seed = np.random.randint(2**31)
    if world_size == 1:
        return seed

    if rank == 0:
        random_num = torch.tensor(seed, dtype=torch.int32, device=device)
    else:
        random_num = torch.tensor(0, dtype=torch.int32, device=device)
    dist.broadcast(random_num, src=0)
    return random_num.item()


def set_random_seed(seed, deterministic=False):
    """Set random seed.

    Args:
        seed (int): Seed to be used.
        deterministic (bool): Whether to set the deterministic option for
            CUDNN backend, i.e., set `torch.backends.cudnn.deterministic`
            to True and `torch.backends.cudnn.benchmark` to False.
            Default: False.
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    if deterministic:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


def train_segmentor(model,
                    dataset,
                    cfg,
                    distributed=False,
                    validate=False,
                    timestamp=None,
                    meta=None):
    """Launch segmentor training."""
    logger = get_root_logger(cfg.log_level)

    # prepare data loaders
    dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
    data_loaders = [
        build_dataloader(
            ds,
            cfg.data.samples_per_gpu,
            cfg.data.workers_per_gpu,
            # cfg.gpus will be ignored if distributed
            len(cfg.gpu_ids),
            dist=distributed,
            seed=cfg.seed,
            drop_last=True) for ds in dataset
    ]
    
    #####
    # added by LYU: 2022/02/23: For Dataloader Vis
    # dataloader_visulization(cfg, data_loaders)
    #####
    
    # put model on gpus
    if distributed:
        find_unused_parameters = cfg.get('find_unused_parameters', False)
        ## modified by LYU: 2022/04/06 
        ## to support gan ddp training
        use_ddp_wrapper = cfg.get('use_ddp_wrapper', False)
        # Sets the `find_unused_parameters` parameter in
        # torch.nn.parallel.DistributedDataParallel
        if use_ddp_wrapper:
            mmcv.print_log('Use DDP Wrapper.', 'mmgen')
            model = DistributedDataParallelWrapper(
                model.cuda(),
                device_ids=[torch.cuda.current_device()],
                broadcast_buffers=False,
                find_unused_parameters=find_unused_parameters)
        else:
            model = MMDistributedDataParallel(
                model.cuda(),
                device_ids=[torch.cuda.current_device()],
                broadcast_buffers=False,
                find_unused_parameters=find_unused_parameters)
    else:
        if not torch.cuda.is_available():
            assert digit_version(mmcv.__version__) >= digit_version('1.4.4'), \
                'Please use MMCV >= 1.4.4 for CPU training!'
        model = MMDataParallel(model, device_ids=cfg.gpu_ids)
    
    # build runner
    ## added by LYU: 2022/03/29
    ## for adapseg model
    if isinstance(cfg.optimizer, dict) and cfg.optimizer.get('backbone', None):
        optimizer = build_optimizers_multi(model, cfg.optimizer)
        if cfg.get('optimizer_cfg', None) is None:
            cfg.optimizer_config = None
        # default to use OptimizerHook
        elif distributed and 'type' not in cfg.optimizer_config:
            cfg.optimizer_config = OptimizerHook(**cfg.optimizer_config)
        else:
            cfg.optimizer_config = cfg.optimizer_config
    ## added by LYU: 2022/04/22:
    ## for DA model
    elif isinstance(cfg.optimizer, dict) and cfg.optimizer.get('backbone_s', None):
        optimizer = build_optimizers_multi(model, cfg.optimizer)
        if cfg.get('optimizer_cfg', None) is None:
            cfg.optimizer_config = None
        # default to use OptimizerHook
        elif distributed and 'type' not in cfg.optimizer_config:
            cfg.optimizer_config = OptimizerHook(**cfg.optimizer_config)
        else:
            cfg.optimizer_config = cfg.optimizer_config
    ## for general segmentation model
    else:
        optimizer = build_optimizer(model, cfg.optimizer)
    
    if cfg.get('runner') is None:
        cfg.runner = {'type': 'IterBasedRunner', 'max_iters': cfg.total_iters}
        warnings.warn(
            'config is now expected to have a `runner` section, '
            'please set `runner` in your config.', UserWarning)
    runner = build_runner(
        cfg.runner,
        default_args=dict(
            model=model,
            batch_processor=None,
            optimizer=optimizer,
            work_dir=cfg.work_dir,
            logger=logger,
            meta=meta))

    # register hooks
    runner.register_training_hooks(cfg.lr_config, cfg.optimizer_config,
                                   cfg.checkpoint_config, cfg.log_config,
                                   cfg.get('momentum_config', None))

    # an ugly walkaround to make the .log and .log.json filenames the same
    runner.timestamp = timestamp

    # register eval hooks
    if validate:
        val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))
        val_dataloader = build_dataloader(
            val_dataset,
            samples_per_gpu=1,
            workers_per_gpu=cfg.data.workers_per_gpu,
            dist=distributed,
            shuffle=False)
        eval_cfg = cfg.get('evaluation', {})
        eval_cfg['by_epoch'] = cfg.runner['type'] != 'IterBasedRunner'
        ## added by LYU: 2022/04/01
        ## For adapseg trainval
        if cfg.runner['type'] == 'DynamicIterBasedRunner':
            eval_cfg['by_epoch'] = False
        eval_hook = DistEvalHook if distributed else EvalHook
        # In this PR (https://github.com/open-mmlab/mmcv/pull/1193), the
        # priority of IterTimerHook has been modified from 'NORMAL' to 'LOW'.
        eval_hook = eval_hook(val_dataloader, **eval_cfg)
        runner.register_hook(eval_hook, priority='LOW')

    if 'ada' in cfg.data:
        ada_dataset = build_dataset(cfg.data.ada, dict(test_mode=True))
        ada_dataloader = build_dataloader(
            ada_dataset,
            samples_per_gpu=1,
            workers_per_gpu=cfg.data.workers_per_gpu,
            dist=distributed,
            shuffle=False)
        ada_cfg = cfg.get('ada_evaluation', {})
        ada_cfg['by_epoch'] = cfg.runner['type'] != 'IterBasedRunner'
        ## added by LYU: 2022/04/01
        ## For adapseg trainval
        if cfg.runner['type'] == 'DynamicIterBasedRunner':
            ada_cfg['by_epoch'] = False
        ada_hook = DistEvalHook if distributed else EvalHook
        # In this PR (https://github.com/open-mmlab/mmcv/pull/1193), the
        # priority of IterTimerHook has been modified from 'NORMAL' to 'LOW'.
        ada_hook = ada_hook(ada_dataloader, **ada_cfg)
        runner.register_hook(ada_hook, priority='LOW')
            # user-defined hooks
    if cfg.get('custom_hooks', None):
        custom_hooks = cfg.custom_hooks
        assert isinstance(custom_hooks, list), \
            f'custom_hooks expect list type, but got {type(custom_hooks)}'
        for hook_cfg in cfg.custom_hooks:
            assert isinstance(hook_cfg, dict), \
                'Each item in custom_hooks expects dict type, but got ' \
                f'{type(hook_cfg)}'
            hook_cfg = hook_cfg.copy()
            priority = hook_cfg.pop('priority', 'NORMAL')
            hook = build_from_cfg(hook_cfg, HOOKS)
            runner.register_hook(hook, priority=priority)

    if cfg.resume_from is None and cfg.get('auto_resume'):
        resume_from = find_latest_checkpoint(cfg.work_dir)
        if resume_from is not None:
            cfg.resume_from = resume_from
    if cfg.resume_from:
        runner.resume(cfg.resume_from)
    elif cfg.load_from:
        runner.load_checkpoint(cfg.load_from)
    
    ## add gan_runner support by LYU: 2022/03/30
    if cfg.get('total_iters', None):
        runner.run(data_loaders, cfg.workflow, cfg.total_iters)
    else:
        runner.run(data_loaders, cfg.workflow)